import os
import torch
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import cv2
from model import AdvancedResNet50, ResNet50WithAttention, ResNet50Model

# 设置matplotlib字体，解决中文显示问题
def set_matplotlib_font():
    """设置matplotlib字体以支持中文显示"""
    import matplotlib.font_manager as fm
    # 尝试多种常见中文字体
    chinese_fonts = ['SimHei', 'Microsoft YaHei', 'SimSun', 'NSimSun', 'FangSong', 'KaiTi']
    font_found = False
    
    for font_name in chinese_fonts:
        # 检查字体是否存在
        font_path = fm.findfont(fm.FontProperties(family=font_name))
        if os.path.exists(font_path) and not font_path.endswith('DejaVuSans.ttf'):
            plt.rcParams['font.sans-serif'] = [font_name] + plt.rcParams['font.sans-serif']
            plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号
            print(f"使用中文字体: {font_name}")
            font_found = True
            break
    
    if not font_found:
        print("警告: 未找到支持中文的字体，图表中的中文可能无法正确显示")
        # 使用英文标签作为备选方案
        plt.rcParams['font.sans-serif'] = ['Arial'] + plt.rcParams['font.sans-serif']

# 应用字体设置
set_matplotlib_font()

# 定义类别名称
CLASS_NAMES = [
    'Chameleon（变色龙）', 
    'Crocodile_Alligator（鳄鱼）', 
    'Frog（青蛙）', 
    'Gecko（壁虎）', 
    'Iguana（鬣蜥）', 
    'Lizard（蜥蜴）', 
    'Salamander（蝾螈）', 
    'Snake（蛇）', 
    'Toad（蟾蜍）', 
    'Turtle_Tortoise（龟/陆龟）'
]

# 定义预处理转换
def get_transforms(image_size=224):
    """获取预测用的数据变换"""
    predict_transform = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return predict_transform

# 加载模型
def load_model(model_path='model/best_model_advanced.pth', model_type='advanced', num_classes=10):
    """加载预训练模型"""
    # 检查GPU可用性
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 创建模型
    if model_type == 'advanced':
        model = AdvancedResNet50(num_classes=num_classes)
    elif model_type == 'attention':
        model = ResNet50WithAttention(num_classes=num_classes)
    else:
        model = ResNet50Model(num_classes=num_classes)
    
    # 加载权重
    if os.path.exists(model_path):
        model.load_state_dict(torch.load(model_path, map_location=device))
        print(f"模型已加载: {model_path}")
    else:
        print(f"警告: 模型文件 {model_path} 不存在！")
    
    model = model.to(device)
    model.eval()
    
    return model, device

# 创建CAM热力图
def generate_cam(model, input_tensor, target_layer_name='layer4'):
    """生成类激活映射(CAM)热力图"""
    # 获取特征图
    features = []
    
    def hook_fn(module, input, output):
        features.append(output.detach())
    
    # 注册钩子
    if hasattr(model.model, target_layer_name):
        target_layer = getattr(model.model, target_layer_name)
        hook = target_layer.register_forward_hook(hook_fn)
    else:
        print(f"警告: 模型中没有 {target_layer_name} 层")
        return None
    
    # 前向传播
    with torch.no_grad():
        output = model(input_tensor)
    
    # 移除钩子
    hook.remove()
    
    # 获取权重
    params = list(model.parameters())
    weight_softmax = params[-2].detach().cpu().numpy()
    
    # 获取预测类别
    _, predicted_class = torch.max(output, 1)
    predicted_class = predicted_class.item()
    
    # 获取特征图
    feature_map = features[0].cpu().numpy().squeeze()
    
    # 生成CAM
    cam = np.zeros(feature_map.shape[1:], dtype=np.float32)
    for i, w in enumerate(weight_softmax[predicted_class]):
        cam += w * feature_map[i, :, :]
    
    # 归一化
    cam = np.maximum(cam, 0)
    cam = cv2.resize(cam, (224, 224))
    cam = cam - np.min(cam)
    cam = cam / np.max(cam) if np.max(cam) > 0 else cam
    
    return cam, predicted_class

# 预测函数
def predict_image(image, model, device, transform):
    """预测图像类别"""
    # 转换图像
    img_tensor = transform(image).unsqueeze(0).to(device)
    
    # 生成CAM热力图
    cam, predicted_class = generate_cam(model, img_tensor)
    
    # 获取类别概率
    with torch.no_grad():
        outputs = model(img_tensor)
        probabilities = torch.nn.functional.softmax(outputs, dim=1)[0].cpu().numpy()
    
    # 创建可视化
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 7))
    
    # 原图
    ax1.imshow(image)
    ax1.set_title('原始图像')
    ax1.axis('off')
    
    # 热力图
    if cam is not None:
        # 创建自定义颜色映射
        colors = [(0, 0, 1), (0, 1, 1), (0, 1, 0), (1, 1, 0), (1, 0, 0)]
        cmap = LinearSegmentedColormap.from_list('custom_cmap', colors, N=256)
        
        # 将热力图叠加到原图上
        img_array = np.array(image)
        img_array = cv2.resize(img_array, (224, 224))
        heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
        
        # 叠加
        alpha = 0.5
        superimposed_img = heatmap * alpha + img_array * (1 - alpha)
        superimposed_img = superimposed_img.astype(np.uint8)
        
        ax2.imshow(superimposed_img)
        ax2.set_title('类激活热力图')
        ax2.axis('off')
    else:
        ax2.imshow(image)
        ax2.set_title('热力图不可用')
        ax2.axis('off')
    
    plt.tight_layout()
    
    # 创建预测结果表格
    results = []
    for i, (label, prob) in enumerate(zip(CLASS_NAMES, probabilities)):
        # 将概率格式化为百分比，保留两位小数
        prob_formatted = f"{prob*100:.2f}%"
        results.append((label, prob_formatted, "✓" if i == predicted_class else ""))
    
    # 按概率排序（需要先转换回浮点数进行排序）
    results.sort(key=lambda x: float(x[1].rstrip('%')), reverse=True)
    
    # 返回结果
    return fig, results

# Gradio界面
def create_interface():
    """创建Gradio界面"""
    # 加载模型
    model, device = load_model()
    transform = get_transforms()
    
    # 定义接口
    def process_image(input_image):
        if input_image is None:
            return None, None
        
        # 预测图像
        fig, results = predict_image(input_image, model, device, transform)
        
        # 创建DataFrame
        headers = ["类别", "概率", "预测结果"]
        
        return fig, results
    
    # 创建界面
    with gr.Blocks(title="爬行动物和两栖动物分类系统", theme=gr.themes.Soft()) as interface:
        gr.Markdown("# 爬行动物和两栖动物分类系统 ")
        gr.Markdown("上传一张爬行动物或两栖动物的图片，系统将自动识别其种类。")
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(type="pil", label="上传图片")
                submit_btn = gr.Button("开始识别", variant="primary")
            
            with gr.Column(scale=2):
                output_plot = gr.Plot(label="可视化结果")
                output_table = gr.Dataframe(headers=["类别", "概率", "预测结果"], label="分类结果")
        
        gr.Markdown("### 系统说明")
        gr.Markdown("""
        - 本系统可以识别10种爬行动物和两栖动物：变色龙、鳄鱼、青蛙、壁虎、鬣蜥、蜥蜴、蝾螈、蛇、蟾蜍、龟/陆龟
        - 左侧图像显示原始上传图片，右侧显示类激活热力图，用于解释模型关注的区域
        - 表格显示每个类别的预测概率，按概率从高到低排序
        """)
        
        # 设置事件
        submit_btn.click(fn=process_image, inputs=input_image, outputs=[output_plot, output_table])
        
    
    return interface

# 主函数
def main():
    """主函数"""
    
    # 创建并启动界面
    interface = create_interface()
    interface.launch(share=True)

if __name__ == "__main__":
    main()